[Bugfix] Padded Eagle Specdec with Chunked Prefill#26263
[Bugfix] Padded Eagle Specdec with Chunked Prefill#26263benchislett merged 11 commits intovllm-project:mainfrom
Conversation
Signed-off-by: Rémi Delacourt <remi@mistral.ai>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request provides a bugfix for padded speculative decoding, specifically within the prepare_next_token_ids_padded function. The change removes a faulty special case for single-token generation (max_gen_len == 1) that did not correctly handle discarded requests. This could have resulted in invalid token IDs being used in subsequent steps. The new implementation generalizes the validity masking logic, making it simpler, more robust, and correct for all scenarios. The fix is sound and improves the correctness of the speculative decoding implementation.
| valid_mask = (valid_sampled_token_ids_gpu != -1) & ( | ||
| valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size | ||
| ) |
There was a problem hiding this comment.
This simplification is a great improvement. By removing the special case for max_gen_len == 1, the code is now more robust. The previous logic didn't account for discarded requests when max_gen_len == 1, which could lead to using an invalid token ID of -1. This unified approach correctly handles all cases.
Flechman
changed the title
|
This bug has been fixed in #26231. I think it would still be nice to merge updated tests, so please update your PR with the fix from main if you wish to continue. |
|
@benchislett sounds good! Done. |
Signed-off-by: Rémi Delacourt <remi@mistral.ai>
…nto fix-eagle-prepare-padded
Head branch was pushed to by a user without write access
|
@benchislett I enforced the test to run on GPUs. It passes locally with 4xH100. |
Signed-off-by: Rémi Delacourt <54138269+Flechman@users.noreply.github.com>
Signed-off-by: Rémi Delacourt <remi@mistral.ai> Signed-off-by: Rémi Delacourt <54138269+Flechman@users.noreply.github.com> Signed-off-by: remi <remi@mistral.ai> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: soaringk <k3vin.zhang@gmail.com>
Signed-off-by: Rémi Delacourt <remi@mistral.ai> Signed-off-by: Rémi Delacourt <54138269+Flechman@users.noreply.github.com> Signed-off-by: remi <remi@mistral.ai> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Signed-off-by: Rémi Delacourt <remi@mistral.ai> Signed-off-by: Rémi Delacourt <54138269+Flechman@users.noreply.github.com> Signed-off-by: remi <remi@mistral.ai> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
2 participants
Purpose
This adds test to #26231
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.